from flask_wtf import CSRFProtect
from flask_wtf.csrf import generate_csrf, CSRFError

# 实例
csrf = CSRFProtect()


class Csrf(object):
    csrf = CSRFProtect()

    def __init__(self, app=None):
        if app:
            self.init_app(app)

    def init_app(self, app):
        self.csrf.init_app(app)

        @app.errorhandler(CSRFError)
        def handle_csrf_error(e):
            return 'error'

        @app.after_request
        def after_request(response):
            csrf_token = generate_csrf()
            response.set_cookie('csrf_token', csrf_token)
            return response
